import sensor, image, time, lcd
import KPU as kpu
import gc
from machine import UART,Timer
from fpioa_manager import fm
from Maix import GPIO
import math
gc.collect()
clock=time.clock()
lcd.init(freq=10000000)
lcd.rotation(2)
sensor.reset()                      # Reset and initialize the sensor. It will
# sensor.set_windowing(224,224)
# sensor.set_auto_exposure(1)         # run automatically, call sensor.run(0) to stop
sensor.set_pixformat(sensor.RGB565) # Set pixel format to RGB565 (or GRAYSCALE)
sensor.set_framesize(sensor.QVGA)   # Set frame size to QVGA (320x240)

# clock = time.clock()                # Create a clock object to track the FPS.
red_threshold   =(45, 70, 19, 48, 2, 37)
roi1=[0,0,320,25]
roi2=[0,100,319,20]
roi3=[0,210,319,25]

pt1=[0,0,0,0]
ptm=[0,0,0,0]
pt2=[0,0,0,0]
ptRes=[0,0,0,0]


task = kpu.load("/sd/m1.kmodel")
anchors = (2.1875, 2.625, 1.78125, 1.5625, 1.90625, 2.09375, 1.5625, 1.8125, 1.375, 1.4375)
a = kpu.init_yolo2(task, 0.6, 0.3, 5, anchors)
labels = ["2", "4", "1", "3", "5", "6", "7", "8"]

fm.register(GPIO.GPIOHS9, fm.fpioa.UART1_RX, force=True)
fm.register(GPIO.GPIOHS10, fm.fpioa.UART1_TX, force=True)

uart = UART(UART.UART1, 115200, read_buf_len=4096)

while(True):
    clock.tick()
    img = sensor.snapshot()

    '''
    *************line detection**************
    '''
    blobs = img.find_blobs([red_threshold],x_stride=10,roi=roi2,merge=True,margin=5)
    if blobs:
        for b in blobs:
            if ptm[2]*ptm[3]<b[2]*b[3]:
                ptm=b[0:4]

    ls=img.find_lines(threshold=2000,theta_margin=45,rho_margin=100)
    for l in ls:
        img.draw_line(l.line(),color=(0,255,0),thickness=5)

    #too large, or too deflected, or no result at all
    stdm= ptm[2]>60 or ptm[0]<30 or ptm[0]+ptm[2]>290 or ptm[0]==0

    if stdm==0:#if middle point is valid
        ptRes=ptm
    else:
        #else find another two points
        blobs = img.find_blobs([red_threshold],x_stride=10,roi=roi1,merge=True,margin=20)
        if blobs:
            for b in blobs:
                img.draw_rectangle(b[0:4])
                if pt1[2]*pt1[3]<b[2]*b[3]:
                    pt1=b[0:4]
        blobs = img.find_blobs([red_threshold],x_stride=10,roi=roi3,merge=True,margin=20)
        if blobs:
            for b in blobs:
                img.draw_rectangle(b[0:4])
                if pt2[2]*pt2[3]<b[2]*b[3]:
                    pt2=b[0:4]
        std1= pt1[2]>50 or pt1[0]==0
        std2= pt2[2]>50 or pt2[0]<30 or pt2[0]+pt2[2]>290 or pt2[0]==0
        if std1+std2==0:
            ptRes[0]=int((pt1[0]+pt2[0])/2)
            ptRes[1]=int((pt1[1]+pt2[1])/2)
            ptRes[2]=pt1[2]
    if ptRes[0]!=0:
        #1-header 2-pointx_hundreds 3-pointx_digits
        data = bytearray([0x21,int(ptRes[0]/100),int(ptRes[0]%100)])
        uart.write(data)
        img.draw_rectangle(ptRes[0],ptRes[1],ptRes[2],25,color=(0,255,0),scale=2,fill=1)

    pt1=[0,0,0,0]
    pt2=[0,0,0,0]
    ptm=[0,0,0,0]
    ptRes=[0,0,0,0]

    '''
    *************number detection**************
    '''
    img2=img.copy(roi=(48,8,224,224))
    img2.pix_to_ai()
    code = kpu.run_yolo2(task, img2)
    del img2
    if code:
        max1=[0,0,0,0]
        max1_id=0
        max2=[0,0,0,0]
        max2_id=0
        for i in code:
            img.draw_rectangle(i.x()+48,i.y()+8,i.w(),i.h(),(0,255,0),2)
            #find the largest two rect
            if max1[2]*max1[3]<=max2[2]*max2[3]:
                if i.w()*i.h()>max1[2]*max1[3]:
                    max1=i.rect()
                    max1_id=int(labels[i.classid()])

            else:
                if i.w()*i.h()>max2[2]*max2[3]:
                    max2=i.rect()
                    max2_id=int(labels[i.classid()])

        for i in code:
            img.draw_rectangle(i.x()+45, i.y()-8,40,10,color=(0,255,0),scale=1.5,fill=1)
            img.draw_string(i.x()+45, i.y()-8, labels[i.classid()]+" "+'%.2f'%i.value(),color=(255,255,255),scale=1.5)
        #1-header 2-left_id 3-right_id
        data = bytearray([0x22,max1_id if max1[0]<max2[0] else max2_id,max2_id if max1[0]<max2[0] else max1_id])
        uart.write(data)

    lcd.display(img)
    print(clock.fps())

    gc.collect()
a = kpu.deinit(task)
